import os
import torch
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
from sklearn.cluster import KMeans, MiniBatchKMeans
from tqdm import tqdm
import argparse
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('ignore', category=UserWarning)

from utils.common import check_and_create_directory, list_all_files, chunk_list

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Process some paths.')
    parser.add_argument('--mode', type=str, required=True, help='Mode (train/val)')
    args = parser.parse_args()

    mode = args.mode
    input_path = '/home/zry/datasets/building/%s/rgbs' % mode
    output_path = '/home/zry/datasets/building/%s/seg' % mode
    tmp_path = '/home/zry/datasets/building/%s/tmp' % mode

    # 所有pth
    paths = list_all_files(tmp_path, '.pth')
    encs = []
    seg_images = []
    
    for i, path in tqdm(enumerate(paths), total=len(paths)):
        f = torch.load(path, "cpu")
        si, enc = f['seg_images'], f['enc']
        encs.append(enc)
        seg_images+=si
    
    # seg_images = torch.cat(seg_images)
    encs = torch.cat(encs)
    encs = encs.detach().numpy()
    
    # 使用K-Means进行聚类
    n_clusters = 5  # 假设我们选择5个聚类
    kmeans = KMeans(n_clusters=n_clusters)
    labels = kmeans.fit_predict(encs)
    
    # 使用TSNE进行降维
    tsne = TSNE(n_components=2)
    reduced_data = tsne.fit_transform(encs)
    
    # 可视化TSNE结果并根据聚类结果着色
    plt.figure(figsize=(10, 6))
    scatter = plt.scatter(reduced_data[:, 0], reduced_data[:, 1], c=labels, marker='o', cmap='viridis')
    plt.colorbar(scatter)
    plt.grid(True)
    
    # 保存图片文件
    plt.savefig('tsne_reduction_with_clusters.png')
    plt.show()

    # 记录每个图片的id和对应的标签
    seg_labels = {
        "seg_images": seg_images,
        "labels": labels
    }
    torch.save(seg_labels, tmp_path + '/seg_labels.npy')
    
    print('You can contine run `python final.main.py`.')